import numpy as np
import pandas as pd
import os
import matplotlib.pyplot as plt
import datetime as dt
from numpy import newaxis
from keras.layers import Dense, Activation, Dropout, LSTM, Flatten
from keras.models import Sequential, load_model
from keras.callbacks import EarlyStopping, ModelCheckpoint
import datetime as dt

dataframe = pd.read_csv(os.path.join('data', 'SensorID=1.csv'))
dataframe.drop(columns=["DataTime", "SensorID", "DID", "sMinute", "sTMinute", "sHour", "sDay", "sWeek"], inplace=True)

# 数据切分
seq_len = 49
delay = 1
data_ = []
for i in range(len(dataframe) - seq_len - delay):
    data_.append(dataframe.iloc[i:i + seq_len + delay])
data_ = np.array([df.values for df in data_])
x = data_[:, :49, :]
y = data_[:, -1, 0]
split_b = int(data_.shape[0] * 0.85)
train_x = x[:split_b]
train_y = y[:split_b]
test_x = x[split_b:]
test_y = y[split_b:]

# 数据标准化
mean = train_x.mean(axis=0)
std = train_x.std(axis=0)
train_x = (train_x - mean) / std
test_x = (test_x - mean) / std


# 模型建立

class Timer():

    def __init__(self):
        self.start_dt = None

    def start(self):
        self.start_dt = dt.datetime.now()

    def stop(self):
        end_dt = dt.datetime.now()
        print('Time taken: %s' % (end_dt - self.start_dt))


# 建立LSTM
model = Sequential()
model.add(LSTM(100, input_shape=(49, 1), return_sequences=True))
model.add(Dropout(0.4))
model.add(LSTM(100, return_sequences=True))
model.add(LSTM(100, return_sequences=False))
model.add(Dropout(0.2))
model.add(Dense(1, activation="linear"))
model.compile(loss="mse", optimizer="adam", metrics=["mae"])
timer = Timer()
timer.start()
epochs = 10
batch_size = 32
save_fname = os.path.join("saved_models", '%s-e%s.h5' % (dt.datetime.now().strftime('%d%m%Y-%H%M%S'), str(epochs)))
callbacks = [
    EarlyStopping(monitor='val_loss', patience=2),
    ModelCheckpoint(filepath=save_fname, monitor='val_loss', save_best_only=True)
]
history = model.fit(
    train_x,
    train_y,
    epochs=epochs,
    batch_size=batch_size,
    callbacks=callbacks,
    validation_data=(test_x, test_y)
)
model.save(save_fname)
timer.stop()


# 预测结果
def predict_point_by_point(data, debug=False):
    if debug == False:
        print('[Model] Predicting Point-by-Point...')
        predicted = model.predict(data)
        predicted = np.reshape(predicted, (predicted.size,))
    else:
        print('[Model] Predicting Point-by-Point...')
        print(np.array(data).shape)
        predicted = model.predict(data)
        print(np.array(predicted).shape)
        predicted = np.reshape(predicted, (predicted.size,))
        print(np.array(predicted).shape)
    return predicted


predictions_pointbypoint = predict_point_by_point(test_x, debug=True)


# 绘图展示

def plot_results(predicted_data, true_data):
    fig = plt.figure(facecolor='white')
    ax = fig.add_subplot(111)
    ax.plot(true_data, label='True Data')
    plt.plot(predicted_data, label='Prediction')
    plt.xlabel("Time")
    plt.ylabel("stress")
    plt.title("Prediction and True Data")
    plt.legend()
    plt.show()


def plot_results_mae(train_loss, val_loss, train_mae, val_mae):
    plt.figure(facecolor='white')
    plt.plot(train_loss, label='train_loss')
    plt.plot(val_loss, label='val_loss')
    plt.plot(train_mae, label='train_mae')
    plt.plot(val_mae, label='val_mae')
    plt.xlabel("Epoch")
    plt.ylabel("Loss/mae")
    plt.title("Prediction and True Data")
    plt.legend()
    plt.show()


plot_results(predictions_pointbypoint, test_y)

# 这里就训练了一次，当然第二张没有图片
plot_results_mae(history.history["loss"], history.history["val_loss"], history.history["mae"],
                 history.history["val_mae"])
